import torch.nn as nn

class UserTower(nn.Module):
    def __init__(self, num_users, embedding_dim):
        super().__init__()
        self.embedding = nn.Embedding(num_users, embedding_dim)
        self.fc = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim),
            nn.ReLU(),
            nn.Linear(embedding_dim, embedding_dim)
        )
    
    def forward(self, user_ids):
        embeds = self.embedding(user_ids)
        return self.fc(embeds)

class ItemTower(nn.Module):
    def __init__(self, num_items, embedding_dim):
        super().__init__()
        self.embedding = nn.Embedding(num_items, embedding_dim)
        self.fc = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim),
            nn.ReLU(),
            nn.Linear(embedding_dim, embedding_dim)
        )
    
    def forward(self, item_ids):
        embeds = self.embedding(item_ids)
        return self.fc(embeds)

class TwoTowerModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.user_tower = UserTower(config.NUM_USERS, config.EMBEDDING_DIM)
        self.item_tower = ItemTower(config.NUM_ITEMS, config.EMBEDDING_DIM)
    
    def forward_user(self, user_ids):
        return self.user_tower(user_ids)
    
    def forward_item(self, item_ids):
        return self.item_tower(item_ids)
    
    def forward(self, user_ids, item_ids):
        user_emb = self.forward_user(user_ids)
        item_emb = self.forward_item(item_ids)
        return (user_emb * item_emb).sum(dim=1)